from json import load
import time
import boto3
from collections import defaultdict, OrderedDict
from time import sleep

from benchmark.utils import Print, BenchError, progress_bar
from ebbflow.settings import Settings, SettingsError

class InstanceManager:
    INSTANCE_NAME = 'EbbFlow-Benchmark'
    SECURITY_GROUP_NAME = 'EbbFlow-Benchmark-SG'
    VPC_NAME = 'EbbFlow-VPC'

    def __init__(self, settings):
        assert isinstance(settings, Settings)
        self.settings = settings
        self.ec2_clients = OrderedDict()
        self.security_groups = OrderedDict()
        
        # 为每个地区创建一个Client
        for region in settings.aws_regions:
            self.ec2_clients[region] = boto3.client('ec2', region_name=region)

    @classmethod
    def make(cls, settings_file='settings.json'):
        try:
            return cls(Settings.load(settings_file))
        except SettingsError as e:
            raise BenchError('Failed to load settings', e)

    def _get(self, state):
        try:
            ids, ips = defaultdict(list), defaultdict(list)
            for region, client in self.ec2_clients.items():
                response = client.describe_instances(
                    Filters=[
                        {'Name': 'tag:Name', 'Values': [self.INSTANCE_NAME]},
                        {'Name': 'instance-state-name', 'Values': state}
                    ]
                )
                
                for reservation in response['Reservations']:
                    for instance in reservation['Instances']:
                        if instance['State']['Name'] in state:
                            ids[region] += [instance['InstanceId']]
                            if instance.get('PublicIpAddress'):
                                ips[region] += [instance['PublicIpAddress']]

        except Exception as error:
            print(f"Error: {error}")

        return ids, ips

    def _wait(self, state):
        while True:
            sleep(1)
            ids, _ = self._get(state)
            if sum(len(x) for x in ids.values()) == 0:
                break

    def _create_security_group(self, client, region):
        try:
            # 获取默认VPC
            vpcs = client.describe_vpcs(Filters=[{'Name': 'is-default', 'Values': ['true']}])
            if not vpcs['Vpcs']:
                raise Exception(f"No default VPC found in {region}")
            vpc_id = vpcs['Vpcs'][0]['VpcId']
            
            # 创建安全组
            response = client.create_security_group(
                GroupName=self.SECURITY_GROUP_NAME,
                Description=f'Security group for {self.INSTANCE_NAME}',
                VpcId=vpc_id
            )
            security_group_id = response['GroupId']
            self.security_groups[region] = security_group_id
            
            # 添加SSH规则
            client.authorize_security_group_ingress(
                GroupId=security_group_id,
                IpPermissions=[
                    {
                        'IpProtocol': 'tcp',
                        'FromPort': 22,
                        'ToPort': 22,
                        'IpRanges': [{'CidrIp': '0.0.0.0/0'}]
                    }
                ]
            )
            
            # 添加EbbFlow端口规则
            client.authorize_security_group_ingress(
                GroupId=security_group_id,
                IpPermissions=[
                    {
                        'IpProtocol': 'tcp',
                        'FromPort': 8000,
                        'ToPort': 8010,
                        'IpRanges': [{'CidrIp': '0.0.0.0/0'}]
                    }
                ]
            )
            
            print(f"✅ 安全组创建成功: {region} - {security_group_id}")
            
        except Exception as e:
            print(f"❌ 安全组创建失败: {region} - {e}")
            raise BenchError('Failed to create security group', e)

    def create_instances(self, instances):
        assert isinstance(instances, int) and instances > 0

        # Create the security group in every region.
        for region, client in self.ec2_clients.items():
            try:
                self._create_security_group(client, region)
            except Exception as e:
                raise BenchError('Failed to create security group', e)

        try:
            # Create all instances.
            size = instances * len(self.ec2_clients)
            progress = progress_bar(
                self.ec2_clients.items(), prefix=f'Creating {size} instances'
            )
            for region, client in progress:
                response = client.run_instances(
                    ImageId='ami-0c02fb55956c7d316',  # Ubuntu 20.04 LTS
                    MinCount=instances,
                    MaxCount=instances,
                    InstanceType=self.settings.instance_type,
                    KeyName=self.settings.key_name,
                    SecurityGroupIds=[self.security_groups[region]],
                    TagSpecifications=[
                        {
                            'ResourceType': 'instance',
                            'Tags': [
                                {'Key': 'Name', 'Value': self.INSTANCE_NAME},
                                {'Key': 'Project', 'Value': 'EbbFlow-Benchmark'}
                            ]
                        }
                    ]
                )

            # Wait for the instances to boot.
            Print.info('Waiting for all instances to boot...')
            self._wait(['pending'])
            Print.heading(f'Successfully created {size} new instances')
        except Exception as error:
            print(f"Error: {error}")
            raise BenchError('Failed to create instances', error)

    def terminate_instances(self):
        try:
            ids, _ = self._get(['pending', 'running', 'stopping', 'stopped'])
            size = sum(len(x) for x in ids.values())
            if size != 0:
                # Terminate instances.
                for region, client in self.ec2_clients.items():
                    if ids[region]:
                        client.terminate_instances(InstanceIds=ids[region])

                # Wait for all instances to properly shut down.
                Print.info('Waiting for all instances to shut down...')
                self._wait(['pending', 'running', 'stopping', 'stopped'])
            Print.heading(f'All instances are shut down')
        except Exception as e:
            raise BenchError('Failed to terminate instances', e)

    def delete_security(self):
        for region, client in self.ec2_clients.items():
            if region in self.security_groups:
                try:
                    client.delete_security_group(GroupId=self.security_groups[region])
                    print(f"✅ 安全组删除成功: {region}")
                except Exception as e:
                    print(f"❌ 安全组删除失败: {region} - {e}")

    def start_instances(self, max):
        size = 0
        try:
            ids, _ = self._get(['stopping', 'stopped'])
            for region, client in self.ec2_clients.items():
                if ids[region]:
                    target = ids[region]
                    target = target if len(target) < max else target[:max]
                    size += len(target)
                    client.start_instances(InstanceIds=target)
            Print.heading(f'Starting {size} instances')
        except Exception as error:
            print(f"Error: {error}")

    def stop_instances(self):
        try:
            ids, _ = self._get(['pending', 'running'])
            for region, client in self.ec2_clients.items():
                if ids[region]:
                    client.stop_instances(InstanceIds=ids[region])
            size = sum(len(x) for x in ids.values())
            Print.heading(f'Stopping {size} instances')
        except Exception as error:
            print(f"Error: {error}")

    def hosts(self, flat=False):
        _, ips = self._get(['pending', 'running'])
        return [x for y in ips.values() for x in y] if flat else ips

    def print_info(self):
        hosts = self.hosts()
        key = self.settings.key_path
        text = ''
        for region, ips in hosts.items():
            text += f'\n Region: {region.upper()}\n'
            for i, ip in enumerate(ips):
                new_line = '\n' if (i+1) % 6 == 0 else ''
                text += f'{new_line} {i}\tssh -i {key} ubuntu@{ip}\n'
        print(
            '\n'
            '----------------------------------------------------------------\n'
            ' EBBFLOW INFO:\n'
            '----------------------------------------------------------------\n'
            f' Available machines: {sum(len(x) for x in hosts.values())}\n'
            f'{text}'
            '----------------------------------------------------------------\n'
        )
